{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8b0e0905",
   "metadata": {},
   "source": [
    "# Image Generation\n",
    "\n",
    "This example will show you how to use GPTCache and OpenAI to implement image generation, i.e. to generate relevant images based on text descriptions. Where the OpenAI model will be used to generate the images, and GPTCache will cache the generated images so that the next time the same or similar text description is requested, it can be returned directly from the cache, which can improve efficiency and reduce costs.\n",
    "\n",
    "\n",
    "This bootcamp is divided into three parts: how to initialize gptcache, running the openai model to generate images, and finally showing how to start the service with gradio. You can also try this example on [Google Colab](https://colab.research.google.com/drive/1YRNR5MLux9UbbWs0SgGI3OSPfkPqIBrM?usp=share_link)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b0f837",
   "metadata": {},
   "source": [
    "## Initialize the gptcache\n",
    "\n",
    "Please [install gptcache](https://gptcache.readthedocs.io/en/latest/index.html#) first, then we can initialize the cache. There are two ways to initialize the cache, the first is to use the map cache (exact match cache) and the second is to use the database cache (similar search cache), it is more recommended to use the second one, but you have to install the related requirements.\n",
    "\n",
    "Before running the example, make sure the `OPENAI_API_KEY` environment variable is set by executing `echo $OPENAI_API_KEY`. If it is not already set, it can be set by using `export OPENAI_API_KEY=YOUR_API_KEY` on Unix/Linux/MacOS systems or `set OPENAI_API_KEY=YOUR_API_KEY` on Windows systems."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46326449",
   "metadata": {},
   "source": [
    "### 1. Init for exact match cache\n",
    "\n",
    "`cache.init` is used to initialize gptcache, the default is to use map to search for cached data, `pre_embedding_func` is used to pre-process the data inserted into the cache, more configuration refer to [initialize Cache](https://gptcache.readthedocs.io/en/latest/references/gptcache.html#module-gptcache.Cache)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bbd4d14f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from gptcache import cache\n",
    "# from gptcache.adapter import openai\n",
    "# from gptcache.processor.pre import get_prompt\n",
    "\n",
    "# cache.init(pre_embedding_func=get_prompt)\n",
    "# cache.set_openai_key()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3ebda74",
   "metadata": {},
   "source": [
    "### 2. Init for similar match cache\n",
    "\n",
    "\n",
    "When initializing gptcahe, the following four parameters are configured:\n",
    "\n",
    "- `pre_embedding_func`: pre-processing before extracting feature vectors\n",
    "- `embedding_func`: the method to extract the text feature vector\n",
    "- `data_manager`: DataManager for cache management\n",
    "- `similarity_evaluation`: the evaluation method after the cache hit\n",
    "\n",
    "The `data_manager` is used to store text, feature vector, and image object data, in the example, it takes [Milvus](https://milvus.io/docs) (please make sure it is started), you can also configure other vector storage, refer to [VectorBase API](https://gptcache.readthedocs.io/en/latest/references/manager.html#module-gptcache.manager.vector_data). Also you can [set ObjectBase](https://gptcache.readthedocs.io/en/latest/references/manager.html#module-gptcache.manager.object_data) to configure which method to use to save the generated image, this example will be stored locally, you can also set it to S3 storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "db7e64ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from gptcache import cache\n",
    "from gptcache.adapter import openai\n",
    "from gptcache.processor.pre import get_prompt\n",
    "\n",
    "from gptcache.embedding import Onnx\n",
    "from gptcache.similarity_evaluation.distance import SearchDistanceEvaluation\n",
    "from gptcache.manager import get_data_manager, CacheBase, VectorBase, ObjectBase\n",
    "\n",
    "\n",
    "onnx = Onnx()\n",
    "cache_base = CacheBase('sqlite')\n",
    "vector_base = VectorBase('milvus', host='localhost', port='19530', dimension=onnx.dimension)\n",
    "object_base = ObjectBase('local', path='./images')\n",
    "data_manager = get_data_manager(cache_base, vector_base, object_base)\n",
    "\n",
    "cache.init(\n",
    "    pre_embedding_func=get_prompt,\n",
    "    embedding_func=onnx.to_embeddings,\n",
    "    data_manager=data_manager,\n",
    "    similarity_evaluation=SearchDistanceEvaluation(),\n",
    "    )\n",
    "cache.set_openai_key()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4201d301",
   "metadata": {},
   "source": [
    "## Run openai image generation \n",
    "\n",
    "Then run `openai.Image.create` to generate the image. The generated images can have a size of \"256x256\", \"512x512\", or \"1024x1024\" pixels, and smaller sizes are faster to generate.\n",
    "\n",
    "\n",
    "Note that `openai` here is imported from `gptcache.adapter.openai`, which can be used to cache with gptcache at request time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7555c46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = openai.Image.create(\n",
    "  prompt=\"a white siamese cat\",\n",
    "  n=1,\n",
    "  size=\"256x256\"\n",
    ")\n",
    "image_url = response['data'][0]['url']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d816d5",
   "metadata": {},
   "source": [
    "## Start with gradio\n",
    "\n",
    "Finally, we can start a gradio application for image generation.\n",
    "\n",
    "First define the `image_generation` method, which is used to generate an image based on the input text and also return whether the cache hit or not. Then start the service with gradio, as shown below:\n",
    "\n",
    "![](../assets/image_generation_gradio.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b160e01f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_generation(prompt):\n",
    "    response = openai.Image.create(\n",
    "      prompt=prompt,\n",
    "      n=1,\n",
    "      size=\"256x256\"\n",
    "    )\n",
    "    return response['data'][0]['url'], response.get(\"gptcache\", False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5f12582",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio\n",
    "\n",
    "interface = gradio.Interface(image_generation, \n",
    "                             gradio.Textbox(lines=1, placeholder=\"Description Here...\"),\n",
    "                             [gradio.Image(shape=(200, 200)), gradio.Textbox(label=\"is hit\")]\n",
    "                            )\n",
    "\n",
    "interface.launch(inline=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b5959d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
